#### Balance ####
library(gregmisc)

## no blocking
n <- 20
r <- n/2
cov.imb <- function(mat){
  inv.mat <- matrix(NA, nrow(mat), ncol(mat))
  for(i in 1:nrow(inv.mat)){
    inv.mat[i,] <- (1:n)[!(1:n %in% mat[i,])]
  }
  vec <- abs(apply(mat, 1, mean) - apply(inv.mat, 1, mean))
  return(vec)
}
mat.no <- combinations(n,r)
vec.no <- cov.imb(mat.no)
summary(vec.no)

## worst blocking
mat.wo <- matrix(NA, 2^r, r)
sws <- 2^((ncol(mat.wo)-1):0)
for(co in 1:ncol(mat.wo)){
  smp <- c(co, n-(co-1))
  xysets <- 2^(co-1)
  sw <- sws[co]
  xlims <- c(1:sw)
  if(xysets>1){
    for(st in 2:xysets){
      xstart <- 2*(st-1)
      newlims <- (xstart*sw+1):((xstart+1)*sw)
      xlims <- c(xlims, newlims)
    }
  }
  mat.wo[xlims,co] <- smp[1]
  mat.wo[is.na(mat.wo[,co]),co] <- smp[2]
}

vec.wo <- cov.imb(mat.wo)
summary(vec.wo)

## best blocking
mat.be <- matrix(NA, 2^r, r)
mat.be <- 2*col(mat.be)    
sws <- 2^((ncol(mat.be)-1):0)
for(co in 1:ncol(mat.be)){
  xysets <- 2^(co-1)
  sw <- sws[co]
  xlims <- c(1:sw)
  if(xysets>1){
    for(st in 2:xysets){
      xstart <- 2*(st-1)
      newlims <- (xstart*sw+1):((xstart+1)*sw)
      xlims <- c(xlims, newlims)
    }
  }
  mat.be[xlims,co] <- 2*co-1
}

vec.be <- cov.imb(mat.be)
summary(vec.be)

mean(vec.no==0)  ##[1] 0.02948754
mean(vec.be==0)  ##[1] 0.2460938
mean(vec.no<1)   ##[1] 0.2606356
mean(vec.be<mean(vec.no)) ##[1] 1

## plot
hist(vec.no, freq = FALSE, xlim =c(0, 10), ylim=c(0,4), col = "yellow", xlab="Covariate Imbalance", main = "")
par(new = TRUE)
hist(vec.be, freq = FALSE, xlim =c(0, 10), ylim=c(0,4), col = "red", xlab = "", ylab = "", main = "")
par(new = TRUE)
hist(vec.wo, freq = FALSE, xlim =c(0, 10), ylim=c(0,4), density = 3, col = "blue", xlab = "", ylab = "", main = "")


#### Efficiency ####

dat <- data.frame(id = 1:6, x = c(2,2,4,3,3,4), y1 = c(1,3,4,9,7,6), y0=c(1,3,4,9,7,6))

## analytic sols: CR
library(gregmisc)
combs <- combinations(6,3)
ests.cra <- array(NA,20)
for(nn in 1:nrow(combs)){
  tr <- combs[nn,]
  co <- (1:6)[!(1:6 %in% tr)]
  dd <- dat$y1[tr]-dat$y0[co]
  ests.cra[nn] <- mean(dd)
}
mean(ests.cra)  ##[1] 0
sd(ests.cra)    ##[1] 2.427908

## analytic sols: Blocking
library(gregmisc)
combs <- combinations(6,3)
idx <- array(TRUE, nrow(combs))
for(i in 1:nrow(combs)){
  if((sum(c(1,2) %in% combs[i,])==2) || (sum(c(3,6) %in% combs[i,])==2) ||
     (sum(c(4,5) %in% combs[i,])==2)){
    idx[i] <- FALSE
  }
}
combs <- combs[idx,]
  
ests.bla <- array(NA, nrow(combs))
for(nn in 1:nrow(combs)){
  tr <- combs[nn,]
  co <- (1:6)[!(1:6 %in% tr)]
  dd <- dat$y1[tr]-dat$y0[co]
  ests.bla[nn] <- mean(dd)
}
mean(ests.bla) ##[1] 0
sd(ests.bla)   ##[1] 1.234427

(sd(ests.cra)-sd(ests.bla))/sd(ests.cra)

####### Additional outlier example  ####
pts <- matrix(c(1,1, 2, 1, 1, 4.1, 2, 4, 1, 100), 5, 2, byrow = TRUE)
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=var(pts[1:4,]))), 3)
##[1] 0.000 1.732 1.760 2.449
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=var(pts))), 3)
##[1] 0.000 2.000 0.078 2.032
pts <- matrix(c(1,1, 2, 1, 1, 4.1, 2, 4, 20, 100), 5, 2, byrow = TRUE)
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=var(pts))), 3)
##[1] 0.000 1.719 1.013 0.745

library(MASS)
pts <- matrix(c(1,1, 2, 1, 1, 4.1, 2, 4, 1, 100), 5, 2, byrow = TRUE)
cov.r <- cov.rob(pts[1:4,], method = "mcd")$cov
cov.r2 <- cov.rob(pts, method = "mcd")$cov
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=cov.r)), 3)
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=cov.r2)), 3)
pts <- matrix(c(1,1, 2, 1, 1, 4.1, 2, 4, 20, 100), 5, 2, byrow = TRUE)
cov.r3 <- cov.rob(pts, method = "mcd")$cov
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=cov.r3)), 3)

pts <- matrix(c(1,1, 2, 1, 1, 4.1, 2, 4, 1, 100), 5, 2, byrow = TRUE)
cov.r <- cov.rob(pts[1:4,], method = "mve")$cov
cov.r2 <- cov.rob(pts, method = "mve")$cov
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=cov.r)), 3)
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=cov.r2)), 3)
pts <- matrix(c(1,1, 2, 1, 1, 4.1, 2, 4, 20, 100), 5, 2, byrow = TRUE)
cov.r3 <- cov.rob(pts, method = "mve")$cov
round(sqrt(mahalanobis(pts[1:4,], pts[1,], cov=cov.r3)), 3)
###### End outlier supplementary example ##########

#### Treatment heterogeneity example ####
halfn <- 6
set.seed(905)
junk <- data.frame(x=c(rnorm(halfn), rnorm(halfn, m=10)))
junk$y0 <- junk$x
junk$y1 <- 2*junk$x
iter <- 1000
dm <- NULL
for(i in 1:iter){
  junk$tcr <- sample(0:1, nrow(junk), rep=TRUE) 
  junk$tra <- sample(c(rep(0, halfn), rep(1, halfn)), nrow(junk), rep=TRUE) ## random allocation
  junk$tbl <- c(sample(0:1), sample(0:1), sample(0:1), sample(0:1), sample(0:1), sample(0:1))
  dm$te <- append(dm$te, c(mean(junk$y1[junk$tcr==1]) - mean(junk$y0[junk$tcr==0]), 
                           mean(junk$y1[junk$tra==1]) - mean(junk$y0[junk$tra==0]),
                           mean(junk$y1[junk$tbl==1]) - mean(junk$y0[junk$tbl==0])))
  dm$meth <- append(dm$meth, c("Complete
                              Random", "Random
                              Allocation", "Blocking"))
}
dm <- data.frame(dm)
true <- mean(junk$y1-junk$y0)

library(ggplot2)
qpl <- qplot(data=dm, x=meth, y=te, geom="boxplot") + geom_hline(y=mean(junk$y1-junk$y0), lty=2, col="grey") + coord_flip() + ylab("Estimated Treatment Effect") + xlab("Method") + scale_x_discrete(limits = rev(levels(dm$meth)))
qpl

aggregate(junk$y1-junk$y0, list(x=junk$x>5), mean)
rm(junk, iter, dm, i, true, halfn, qpl)
